#include <rclcpp/rclcpp.hpp>
#include <rclcpp/qos.hpp>  // 添加此行以包含QoS设置
#include <sensor_msgs/msg/laser_scan.hpp>
#include <opencv2/opencv.hpp>
#include <opencv2/core.hpp>
#include <cmath>
#include <vector>
#include <sophus/se2.hpp>  
#include <Eigen/Core>
#include <Eigen/Dense>
#include <Eigen/Geometry>

using SE2 = Sophus::SE2d;
using Vec2d = Eigen::Vector2d;
using Vec2i = Eigen::Vector2i;
using Vec3b = Eigen::Matrix<char, 3, 1>;
using Vec3d = Eigen::Vector3d;
using Mat3d = Eigen::Matrix3d;
using SO2 = Sophus::SO2d;


class SimpleICPNode : public rclcpp::Node {
public:
    SimpleICPNode() : Node("simple_icp_node") {
        rclcpp::QoS qos(rclcpp::KeepLast(10)); 
        qos.reliability(RMW_QOS_POLICY_RELIABILITY_BEST_EFFORT); // 显式设置可靠传输的QoS策略
        scan_sub_ = this->create_subscription<sensor_msgs::msg::LaserScan>(
            "/scan", qos, std::bind(&SimpleICPNode::laserScanCallback,  this, std::placeholders::_1)); // 订阅激光雷达的 /scan 话题
        this->declare_parameter<std::string>("scan_topic", "/scan");
        scan_topic_ = this->get_parameter("scan_topic").as_string();
        buildmodel();
    }

    // 似然场的模板
    struct ModelPoint {
        ModelPoint(int dx, int dy, float res) : dx_(dx), dy_(dy), residual_(res) {}
        int dx_ = 0;
        int dy_ = 0;
        float residual_ = 0;
        
    };


private:
    SE2 pose_;  // T_W_S
    // Scan2d::Ptr target_ = nullptr;
    // 定义激光雷达的消息类型
    sensor_msgs::msg::LaserScan::SharedPtr source_ = nullptr;
    sensor_msgs::msg::LaserScan::SharedPtr last_scan = nullptr;
    sensor_msgs::msg::LaserScan::SharedPtr current_scan = nullptr;
    std::vector<ModelPoint> model_;  // 2D 模板
    cv::Mat field_;                  // 场函数
    bool has_outside_pts_ = false;   // 是否含有出了这个场的点
    float resolution_ = 60;  // 图像的分辨率（1米中包含的像素点）
    int image_size = 1000;

    void buildmodel() {
        const int range = 20;  // 生成多少个像素的模板
        for (int x = -range; x <= range; ++x) {
            for (int y = -range; y <= range; ++y) {
                model_.emplace_back(x, y, std::sqrt((x * x) + (y * y)));
            }
        }
    }

    void setTargetScan(const sensor_msgs::msg::LaserScan::SharedPtr scan) {
        // target_ = scan;

        // 在target点上生成场函数，创建了一个 1000x1000 的单精度浮点数矩阵，并且矩阵中每个元素的初始值都是 30.0。
        field_ = cv::Mat(1000, 1000, CV_32F, 30.0);

        // 遍历激光扫描的每一个数据点
        for (size_t i = 0; i < scan->ranges.size(); ++i) {
            // 检查扫描范围是否有效
            if (scan->ranges[i] < scan->range_min || scan->ranges[i] > scan->range_max) {
                continue; // 如果无效，跳过该点
            }

            // 计算实际的激光角度
            double real_angle = scan->angle_min + i * scan->angle_increment;
            // 计算(x, y)坐标
            double x = scan->ranges[i] * std::cos(real_angle) ;
            double y = scan->ranges[i] * std::sin(real_angle) ;
            // 计算图像坐标
            int image_x = int(x * resolution_ + image_size / 2);
            int image_y = int(y * resolution_ + image_size / 2);

            // 在(x,y)附近填入场函数
            for (auto& model_pt : model_) {
                int xx = int(image_x  + model_pt.dx_);
                int yy = int(image_y + model_pt.dy_);
                if (xx >= 0 && xx < field_.cols && yy >= 0 && yy < field_.rows &&
                    field_.at<float>(yy, xx) > model_pt.residual_) {
                        field_.at<float>(yy, xx) = model_pt.residual_;
                    }
            }
        }
    }

    void setsourcescan(const sensor_msgs::msg::LaserScan::SharedPtr scan) { source_ = scan; }

    bool AlignGaussNewton(SE2& init_pose) {
        int iterations = 10;
        double cost = 0, lastCost = 0;
        SE2 current_pose = init_pose;
        const int min_effect_pts = 20;  // 最小有效点数
        const int image_boarder = 20;   // 预留图像边界
        // float resolution_ = 100;  // 图像的分辨率（1米中包含的像素点）

        has_outside_pts_ = false;
        for (int iter = 0; iter < iterations; ++iter) {
            Mat3d H = Mat3d::Zero();
            Vec3d b = Vec3d::Zero();
            cost = 0;

            int effective_num = 0;  // 有效点数

            // 遍历source
            for (size_t i = 0; i < source_->ranges.size(); ++i) {
                float r = source_->ranges[i];
                if (r < source_->range_min || r > source_->range_max) {
                    continue;
                }

                float angle = source_->angle_min + i * source_->angle_increment;
                if (angle < source_->angle_min + 30 * M_PI / 180.0 || angle > source_->angle_max - 30 * M_PI / 180.0) {
                    continue;
                }

                float theta = current_pose.so2().log();
                Vec2d pw = current_pose * Vec2d(r * std::cos(angle), r * std::sin(angle));

                // 在field中的图像坐标
                Vec2i pf = (pw * resolution_ + Vec2d(500, 500)).cast<int>();

                if (pf[0] >= image_boarder && pf[0] < field_.cols - image_boarder && pf[1] >= image_boarder &&
                    pf[1] < field_.rows - image_boarder) {
                    effective_num++;

                    // 图像梯度
                    float dx = 0.5 * (field_.at<float>(pf[1], pf[0] + 1) - field_.at<float>(pf[1], pf[0] - 1));
                    float dy = 0.5 * (field_.at<float>(pf[1] + 1, pf[0]) - field_.at<float>(pf[1] - 1, pf[0]));

                    Vec3d J;
                    J << resolution_ * dx, resolution_ * dy,
                        -resolution_ * dx * r * std::sin(angle + theta) + resolution_ * dy * r * std::cos(angle + theta);
                    H += J * J.transpose();

                    float e = field_.at<float>(pf[1], pf[0]);
                    b += -J * e;

                    cost += e * e;
                } else {
                    has_outside_pts_ = true;
                }
            }

            if (effective_num < min_effect_pts) {
                return false;
            }

            // solve for dx
            Vec3d dx = H.ldlt().solve(b);
            if (isnan(dx[0])) {
                break;
            }

            cost /= effective_num;
            if (iter > 0 && cost >= lastCost) {
                break;
            }

            // 输出信息
            RCLCPP_INFO(this->get_logger(), "iter = %f, cost = %f, effect num = %f", iter, cost, effective_num);

            current_pose.translation() += dx.head<2>();
            current_pose.so2() = current_pose.so2() * SO2::exp(dx[2]);
            lastCost = cost;
        }

        init_pose = current_pose;
        return true;
    }
    
    void Visualize2DScan(const sensor_msgs::msg::LaserScan::SharedPtr scan, const SE2& pose, cv::Mat& image, const Vec3b& color, int image_size = 1000,
                            float resolution = 60.0, const SE2& pose_submap = SE2()) {
        if (image.data == nullptr) {
            image = cv::Mat(image_size, image_size, CV_8UC3, cv::Vec3b(255, 255, 255));
        }

        for (size_t i = 0; i < scan->ranges.size(); ++i) {
            if (scan->ranges[i] < scan->range_min || scan->ranges[i] > scan->range_max) {
                continue;
            }

            double real_angle = scan->angle_min + i * scan->angle_increment;
            double x = scan->ranges[i] * std::cos(real_angle);
            double y = scan->ranges[i] * std::sin(real_angle);

            if (real_angle < scan->angle_min + 30 * M_PI / 180.0 || real_angle > scan->angle_max - 30 * M_PI / 180.0) {
                continue;
            }

            Vec2d psubmap = pose_submap.inverse() * (pose * Vec2d(x, y));

            int image_x = int(psubmap[0] * resolution + image_size / 2);
            int image_y = int(psubmap[1] * resolution + image_size / 2);
            if (image_x >= 0 && image_x < image.cols && image_y >= 0 && image_y < image.rows) {
                image.at<Vec3b>(image_y, image_x) = Vec3b(color[0], color[1], color[2]);
            }
        }

        // 同时画出pose自身所在位置
        Vec2d pose_in_image =
            pose_submap.inverse() * (pose.translation()) * double(resolution) + Vec2d(image_size / 2, image_size / 2);
        cv::circle(image, cv::Point2f(pose_in_image[0], pose_in_image[1]), 5, cv::Scalar(color[0], color[1], color[2]), 2);
    }

    cv::Mat GetFieldImage() {
        cv::Mat image(field_.rows, field_.cols, CV_8UC3);
        for (int x = 0; x < field_.cols; ++x) {
            for (int y = 0; y < field_.rows; ++y) {
                float r = field_.at<float>(y, x) * 255.0 / 30.0;
                image.at<Vec3b>(y, x) = Vec3b(uchar(r), uchar(r), uchar(r));
            }
        }
        return image;
    }

    void laserScanCallback(const sensor_msgs::msg::LaserScan::SharedPtr scan_msg) {

        current_scan = scan_msg;
        SE2 pose;

        if (last_scan == nullptr) {
            last_scan = current_scan;
            return;
        }

        setTargetScan(last_scan);
        setsourcescan(current_scan);

        // 高斯牛顿优化
        AlignGaussNewton(pose);

        // 输出对齐结果
        RCLCPP_INFO(this->get_logger(), "Alignment Result: x = %f, y = %f, theta = %f", pose.translation().x(), pose.translation().y(), pose.so2().log());
        
        // 可视化
        cv::Mat image;
        Visualize2DScan(last_scan,SE2(),image,Vec3b(255,0,0)); //target 是蓝色
        Visualize2DScan(current_scan,pose,image,Vec3b(0,0,255)); // source 是红色
        cv::imshow("scan", image);

        // 画出target和它的场函数
        cv::Mat field_image = GetFieldImage();
        Visualize2DScan(last_scan,SE2(),field_image,Vec3b(255,0,0),1000); //target 是蓝色
        cv::imshow("field", field_image);
        cv::waitKey(10);
        // Update the last scan
        last_scan = current_scan;
    }


    std::string scan_topic_;
    rclcpp::Subscription<sensor_msgs::msg::LaserScan>::SharedPtr scan_sub_;
};

int main(int argc, char **argv) {
    rclcpp::init(argc, argv);
    rclcpp::spin(std::make_shared<SimpleICPNode>());
    rclcpp::shutdown();
    return 0;
}
